# -*- coding: utf-8 -*-
"""
Created on Sun Apr 23 10:01:39 2023

@author: lenovo
"""

import xarray as xr
import os
import numpy as np
import math
import scipy.interpolate
from osgeo import gdal
import re
import glob
# import shutil
from tqdm import tqdm # 运行计时

class para:
    climate_path=r'E:\drivedata\era5data'
    dem_path=r'E:\GDEM V3'
    slop_path=r'E:\slop30m'

    landtypedir = r'F:\公司项目\陕西\商洛\land30m'
    out_path=r'F:\公司项目\陕西\商洛\data\GEP'
    inputyear=2023
    lon1=-1
    lon2=-1
    lat1=-1
    lat2=-1
    imlat=0
    imlon=0
    imwidth=3601
    imheight=3601
    num='0'
    # soilr=1.305 #土壤容重，t/m3
################################################
def read_img(filename):
    dataset = gdal.Open(filename)  # 打开文件
    if dataset == None:
        print(filename+"文件无法打开")
        return
    im_width = dataset.RasterXSize  # 栅格矩阵的列数
    im_height = dataset.RasterYSize  # 栅格矩阵的行数
    im_bands = dataset.RasterCount #波段数
    im_geotrans = dataset.GetGeoTransform()  # 仿射矩阵
    im_proj = dataset.GetProjection()  # 地图投影信息
    im_data = dataset.ReadAsArray(0, 0, im_width, im_height).astype(np.float32)  # 将数据写成数组，对应栅格矩阵
    im_lon=[im_geotrans[0]+i*im_geotrans[1] for i in range(im_width)]
    im_lat=[im_geotrans[3]+i*im_geotrans[5] for i in range(im_height)]
    
    del dataset  # 关闭对象，文件dataset   
    return im_data,im_width,im_height,im_bands,im_geotrans,im_proj,im_lon,im_lat
#=========================
# 保存tif文件函数
def write_img(im_data, im_geotrans, im_proj, path, nodata=None):
    if 'int8' in im_data.dtype.name:
        datatype = gdal.GDT_Byte
    elif 'int16' in im_data.dtype.name:
        datatype = gdal.GDT_UInt16
    else:
        datatype = gdal.GDT_Float32
    if len(im_data.shape) == 3:
        im_bands, im_height, im_width = im_data.shape
    elif len(im_data.shape) == 2:
        im_data = np.array([im_data])
        im_bands, im_height, im_width = im_data.shape
    # 创建文件
    driver = gdal.GetDriverByName("GTiff")
    dataset = driver.Create(path, int(im_width), int(im_height), int(im_bands), datatype)
    if (dataset != None):
        dataset.SetGeoTransform(im_geotrans)  # 写入仿射变换参数
        dataset.SetProjection(im_proj)  # 写入投影        
    for i in range(im_bands):
        if (nodata != None):
            dataset.GetRasterBand(i + 1).SetNoDataValue(nodata)
        dataset.GetRasterBand(i + 1).WriteArray(im_data[i])
    del dataset

#=======================   
def readasc(fname):
    res=np.loadtxt(fname,skiprows=6)
    res[res==-9999]=np.nan
    return res
#======================================
def interpolation(lat,lon,data,met):
                
    latc=para.imlat
    lonc=para.imlon
    
    lon_c,lat_c=np.meshgrid(lonc,latc)

    lon_b,lat_b=np.meshgrid(lon,lat)
    #method{‘linear’, ‘nearest’, ‘cubic’}, optional
    res = scipy.interpolate.griddata((lon_b.ravel(), lat_b.ravel()),
                                    data.ravel(), (lon_c, lat_c), 
                                    method=met,fill_value=np.nan)
    return res
#==========================
def get_rc(tp):
    if tp==101: #常绿阔叶
        rc=2.67
    elif tp==102: 
        rc=1.33
    elif tp==103:
        rc=3.02
    elif tp==104:
        rc=0.88
    elif tp==105:
        rc=2.29
    elif tp==61:
        rc=19.2
    elif tp==106:
        rc=4.26
    elif tp==107:
        rc=4.17
    elif tp==108:
        rc=4.17
    elif tp==62:
        rc=19.2
    elif tp==21:
        rc=8.02
    elif tp==22:
        rc=4.78
    elif tp==23:
        rc=9.37
    elif tp==63:
        rc=18.27
    elif tp==41:
        rc=34.7
    elif tp==42:
        rc=46.96
    elif tp==109:
        rc=9.57
    elif tp==110:
        rc=7.9
    elif tp==111:
        rc=19.2
    elif tp==112:
        rc=19.2
    elif tp==24:
        rc=18.27    
    elif tp>=31 and tp<=37: #河流湿地
        rc=0
    elif tp>=51 and tp<=54: 
        rc=45
    elif tp>=64 and tp<=68: 
        rc=19.2      
    else:
        rc=0
    
    rc=rc/100
    return rc
        
def get_pe():
    path=para.climate_path


    # 提取指定经纬度范围内的数据
    lon_range = slice(para.lon1, para.lon2)  # 经度范围
    lat_range = slice(para.lat1, para.lat2)  # 纬度范围

    fn=path+'/Total_precipitation/total_precipitation'+str(para.inputyear)+'.nc'
    da=xr.open_dataset(fn)
    data = da.sel(longitude=lon_range, latitude=lat_range)

    lat=data.latitude.values
    lon=data.longitude.values
    pre=data.tp.values*1000 #m -> mm    
    ypre=np.nansum(pre,axis=0)
    
    fn=path+'/Evaporation/evaporation'+str(para.inputyear)+'.nc'
    da=xr.open_dataset(fn)
    data = da.sel(longitude=lon_range, latitude=lat_range)
    
    ev=data.e.values*1000 #m -> mm    
    yev=np.nansum(ev,axis=0)*(-1)
    
    #插值到30m网格
    met='cubic'
    P = interpolation(lat,lon,ypre,met)
    ET = interpolation(lat,lon,yev,met)
    
    return P,ET
    

################################################    
tiffile_all = glob.glob(para.landtypedir+'\\'+'*.tif')
for fn in tqdm(tiffile_all):
    tif_file_folder,tif_file_name = os.path.split(fn)
    para.num=re.findall('\d+', tif_file_name)
    
    im_data,im_width,im_height,im_bands,im_geotrans,im_proj,im_lon,im_lat=read_img(fn)

    para.imlat=im_lat
    para.imlon=im_lon
    para.imheight=im_height
    para.imwidth=im_width
    para.lat1=im_lat[0]+0.5
    para.lat2=im_lat[-1]-0.5
    para.lon1=im_lon[0]-0.5
    para.lon2=im_lon[-1]+0.5


    p,et=get_pe()

    cols=im_width
    rows=im_height
    Qwr=np.zeros((rows,cols))
    for i in range(0,rows):
        for j in range(0,cols):
            if im_data[i,j]==255:
                Qwr[i,j]=np.nan
            else:
                rc=get_rc(im_data[i,j])
                Qwr[i,j]=(p[i,j]-p[i,j]*rc-et[i,j])/1000
    Qwr[Qwr<0]=0
    
    Qwr[np.isnan(Qwr)]=255            
    fout='H:\\company\\公司项目\\陕西\\商洛\\data\\test'+'\\水源涵养实物量_'+tif_file_name
    write_img(Qwr, im_geotrans, im_proj, fout,255)

    # Cwr=6.58 #水库运营成本
    # Pwe=30.36 #单位库容工程造价
    # Dr=0.02 #水库年折旧率
    
    Cwr=1.2 #水库运营成本
    Pwe=16.4 #单位库容工程造价
    Dr=0.06 #水库年折旧率
    Vwr=Qwr*(Cwr+Pwe*Dr)
    
    tif_file_name=tif_file_name.split('_')[1]
    
    Vwr[np.isnan(Vwr)]=-9999
    fout=para.out_path+'\\水源涵养_'+tif_file_name
    write_img(Vwr, im_geotrans, im_proj, fout,-9999)
    
